import sys
import ctypes
from ctypes import wintypes
from PyQt5.QtWidgets import QApplication, QWidget, QDesktopWidget
from PyQt5.QtCore import Qt, QRect, QPoint
from PyQt5.QtGui import QPainter, QColor, QPen, QCursor

HANDLE_SIZE = 12      # 显示用的手柄大小
HITBOX_SIZE = 25      # 判定范围扩大到 20x20
DRAG_THRESHOLD = 6    # 移动超过该距离后认为是拖拽
BORDER_WIDTH= 2       # 边框宽度

# 应用窗口抓取相关
GW_HWNDNEXT = 2
DWMWA_CLOAKED = 14
SPI_GETWORKAREA = 48

# 状态定义
STATE_PREVIEW = "preview"      # 预览状态：鼠标移动时自动检测窗口
STATE_DRAGGING = "dragging"    # 拖拽创建状态：从零创建选框
STATE_SELECTED = "selected"     # 已选择状态：选框已确定，可调整大小
STATE_RESIZING = "resizing"    # 调整大小状态：正在拖动手柄

user32 = ctypes.windll.user32
try:
    dwmapi = ctypes.windll.dwmapi
except OSError:
    dwmapi = None


class ScreenSelector(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle('屏幕选区')
        self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint | Qt.Tool)
        self.setWindowState(Qt.WindowFullScreen)
        self.setAttribute(Qt.WA_TranslucentBackground)
        self.setMouseTracking(True)
        # 设置焦点策略，确保窗口可以接收焦点
        self.setFocusPolicy(Qt.StrongFocus)

        # 状态机
        self.state = STATE_PREVIEW

        # 选框相关
        self.start = None
        self.end = None

        # 调整大小相关
        self.handle_index = -1
        self.offset = QPoint()

        # 鼠标事件相关
        self.mouse_pressed = False
        self.press_pos = QPoint()

        desktop = QApplication.desktop()
        screen_rect = desktop.screenGeometry()
        self.screen_rect = QRect(0, 0, screen_rect.width(), screen_rect.height())
        self.selection_rect = self.screen_rect


    @staticmethod
    def exec(default = 'screen'):
        app = QApplication(sys.argv)
        selector = ScreenSelector()
        selector.show()
        app.exec_()

        if selector.selection_rect:
            x = selector.selection_rect.x()
            y = selector.selection_rect.y()
            w = selector.selection_rect.width()
            h = selector.selection_rect.height()
            print(f"选中区域: x={x}, y={y}, w={w}, h={h}")
            return [x, y, w, h]
        else:
            print("取消选中区域")
            if default == 'screen':
                return [selector.screen_rect.x(), selector.screen_rect.y(), selector.screen_rect.width(), selector.screen_rect.height()]
            else:
                return None


    def _find_window_rect_at(self, global_pos):
        x = global_pos.x()
        y = global_pos.y()
        target_rect = None

        def is_window_cloaked(hwnd):
            if not dwmapi:
                return False
            cloaked = wintypes.DWORD()
            if dwmapi.DwmGetWindowAttribute(hwnd, DWMWA_CLOAKED, ctypes.byref(cloaked), ctypes.sizeof(cloaked)) != 0:
                return False
            return cloaked.value != 0

        def iter_z_order_windows():
            hwnd = user32.GetTopWindow(None)
            while hwnd:
                yield hwnd
                hwnd = user32.GetWindow(hwnd, GW_HWNDNEXT)

        for hwnd in iter_z_order_windows():
            if hwnd == int(self.winId()):
                continue
            if not user32.IsWindowVisible(hwnd):
                continue
            if user32.IsIconic(hwnd):
                continue
            if is_window_cloaked(hwnd):
                continue

            rect = wintypes.RECT()
            if not user32.GetWindowRect(hwnd, ctypes.byref(rect)):
                continue
            if rect.left == rect.right or rect.top == rect.bottom:
                continue
            if rect.left <= x < rect.right and rect.top <= y < rect.bottom:
                target_rect = QRect(rect.left, rect.top, rect.right - rect.left, rect.bottom - rect.top)
                break

        # 查找完成后，再次确保窗口保持焦点
        if not self.hasFocus():
            self.activateWindow()
            self.raise_()
            self.setFocus()

        if target_rect:
            # 分别检测target_rect的4个边界，不应超过屏幕边界
            if target_rect.left() < 0:
                target_rect.setLeft(0)
            if target_rect.top() < 0:
                target_rect.setTop(0)
            if target_rect.right() > self.screen_rect.right():
                target_rect.setRight(self.screen_rect.right())
            if target_rect.bottom() > self.screen_rect.bottom():
                target_rect.setBottom(self.screen_rect.bottom())
            return target_rect
        else:
            # SystemParametersInfo 用于获取系统参数，SPI_GETWORKAREA 获取工作区域
            rect = wintypes.RECT()
            if user32.SystemParametersInfoW(SPI_GETWORKAREA, 0, ctypes.byref(rect), 0):
                return QRect(rect.left, rect.top, rect.right - rect.left, rect.bottom - rect.top)
            return self.screen_rect


    def _get_handle_region(self, size):
        """根据 size 返回8个手柄矩形区域"""
        if not self.selection_rect:
            return []
        r = self.selection_rect
        s = size
        return [
            QRect(r.left() - s//2, r.top() - s//2, s, s),        # 0 左上
            QRect(r.center().x() - s//2, r.top() - s//2, s, s),  # 1 上中
            QRect(r.right() - s//2, r.top() - s//2, s, s),       # 2 右上
            QRect(r.right() - s//2, r.center().y() - s//2, s, s),# 3 右中
            QRect(r.right() - s//2, r.bottom() - s//2, s, s),    # 4 右下
            QRect(r.center().x() - s//2, r.bottom() - s//2, s, s),#5 下中
            QRect(r.left() - s//2, r.bottom() - s//2, s, s),     # 6 左下
            QRect(r.left() - s//2, r.center().y() - s//2, s, s)  # 7 左中
        ]

    
    def _update_cursor(self, pos):
        """根据当前状态和鼠标位置更新光标样式"""
        cursor_for_handle = [
            Qt.SizeFDiagCursor,  # 左上
            Qt.SizeVerCursor,    # 上中
            Qt.SizeBDiagCursor,  # 右上
            Qt.SizeHorCursor,    # 右中
            Qt.SizeFDiagCursor,  # 右下
            Qt.SizeVerCursor,    # 下中
            Qt.SizeBDiagCursor,  # 左下
            Qt.SizeHorCursor     # 左中
        ]
        if self.state == STATE_PREVIEW:
            self.setCursor(Qt.CrossCursor)
        elif self.state == STATE_DRAGGING:
            self.setCursor(Qt.ArrowCursor)
        elif self.state == STATE_SELECTED:
            for i, handle in enumerate(self._get_handle_region(HITBOX_SIZE)):
                if handle.contains(pos):
                    self.setCursor(QCursor(cursor_for_handle[i]))
                    break
            else:
                # 如果循环正常结束，没有触发break，则认为没有悬停在手柄上
                self.setCursor(Qt.CrossCursor)
        elif self.state == STATE_RESIZING:
            # 调整大小时保持手柄光标
            if self.handle_index >= 0:
                self.setCursor(QCursor(cursor_for_handle[self.handle_index]))


    def paintEvent(self, event):
        painter = QPainter(self)
        painter.fillRect(self.rect(), QColor(0, 0, 0, 150))  # 灰色半透明遮罩

        if self.selection_rect:
            # 显示原始区域
            painter.setCompositionMode(QPainter.CompositionMode_Source)
            painter.fillRect(self.selection_rect, QColor(0, 0, 0, 1))

            # 边框（向内收缩 BORDER_WIDTH / 2，使边框绘制在选框内部）
            border_rect = self.selection_rect.adjusted(
                BORDER_WIDTH // 2, 
                BORDER_WIDTH // 2, 
                -BORDER_WIDTH // 2, 
                -BORDER_WIDTH // 2
            )
            pen = QPen(Qt.white, BORDER_WIDTH, Qt.SolidLine)
            painter.setPen(pen)
            painter.drawRect(border_rect)

            # 拖动手柄（只在已选择或调整大小状态显示）
            if self.state in (STATE_SELECTED, STATE_RESIZING):
                for handle in self._get_handle_region(HANDLE_SIZE):
                    painter.fillRect(handle, QColor(255, 255, 255))


    def mousePressEvent(self, event):
        if event.button() != Qt.LeftButton:
            return

        pos = event.pos()
        self.mouse_pressed = True
        self.press_pos = pos

        # 已选择状态：检查是否点击在手柄上
        if self.state == STATE_SELECTED:
            for i, handle in enumerate(self._get_handle_region(HITBOX_SIZE)):
                if handle.contains(pos):
                    # print("STATE_SELECTED to STATE_RESIZING")
                    self.state = STATE_RESIZING
                    self.handle_index = i
                    self.offset = pos
                    return


    def mouseMoveEvent(self, event):
        pos = event.pos()
        global_pos = self.mapToGlobal(pos)

        # 预览状态：自动检测窗口并更新预览选框
        if self.state == STATE_PREVIEW:
            if not self.mouse_pressed:
                new_rect = self._find_window_rect_at(global_pos)
                if new_rect != self.selection_rect:
                    self.selection_rect = new_rect
                    self.update()
                pass
            else:
                # 预览状态下按下鼠标，检查是否开始拖拽
                if (pos - self.press_pos).manhattanLength() >= DRAG_THRESHOLD:
                    # 切换到拖拽创建状态
                    self.state = STATE_DRAGGING
                    # print("STATE_PREVIEW to STATE_DRAGGING")
                    self.start = self.press_pos
                    self.end = pos
                    self.selection_rect = QRect(self.start, self.end).normalized()
                    self.update()
                    return

        # 拖拽创建状态：更新选框
        elif self.state == STATE_DRAGGING:
            self.end = pos
            self.selection_rect = QRect(self.start, self.end).normalized()
            self.update()

        # 调整大小状态：根据手柄调整选框
        elif self.state == STATE_RESIZING:
            if not self.selection_rect:
                return
            dx = pos.x() - self.offset.x()
            dy = pos.y() - self.offset.y()
            r = self.selection_rect
            if self.handle_index == 0:
                r.setTopLeft(r.topLeft() + QPoint(dx, dy))
            elif self.handle_index == 1:
                r.setTop(r.top() + dy)
            elif self.handle_index == 2:
                r.setTopRight(r.topRight() + QPoint(dx, dy))
            elif self.handle_index == 3:
                r.setRight(r.right() + dx)
            elif self.handle_index == 4:
                r.setBottomRight(r.bottomRight() + QPoint(dx, dy))
            elif self.handle_index == 5:
                r.setBottom(r.bottom() + dy)
            elif self.handle_index == 6:
                r.setBottomLeft(r.bottomLeft() + QPoint(dx, dy))
            elif self.handle_index == 7:
                r.setLeft(r.left() + dx)
            self.selection_rect = r.normalized()
            self.offset = pos
            self.update()

        # 更新鼠标样式
        self._update_cursor(pos)


    def mouseReleaseEvent(self, event):
        if event.button() != Qt.LeftButton:
            return

        # 预览状态：如果已有预览选框，切换到已选择状态
        if self.state == STATE_PREVIEW:
            if self.mouse_pressed and self.selection_rect:
                # print("STATE_PREVIEW to STATE_SELECTED")
                self.state = STATE_SELECTED
                self.update()

        # 拖拽创建状态：切换到已选择状态
        elif self.state == STATE_DRAGGING:
            if self.selection_rect:
                # print("STATE_DRAGGING to STATE_SELECTED")
                self.state = STATE_SELECTED
                self.update()

        # 调整大小状态：切换回已选择状态
        elif self.state == STATE_RESIZING:
            # print("STATE_RESIZING to STATE_SELECTED")
            self.state = STATE_SELECTED
            self.handle_index = -1

        # 重置鼠标状态
        self.mouse_pressed = False


    def mouseDoubleClickEvent(self, event):
        if self.selection_rect:
            x = self.selection_rect.x()
            y = self.selection_rect.y()
            w = self.selection_rect.width()
            h = self.selection_rect.height()
            print(f"选中区域: x={x}, y={y}, w={w}, h={h}")
            self.close()
            QApplication.quit()  # 退出主事件循环，彻底结束程序


    def keyPressEvent(self, event):
        if event.key() in (Qt.Key_Return, Qt.Key_Enter):  # 支持主键盘和小键盘的 Enter
            if self.selection_rect:
                self.close()
                QApplication.quit()
        elif event.key() == Qt.Key_Escape:
            self.selection_rect = None
            self.close()
            QApplication.quit()

    

if __name__ == '__main__':
    ScreenSelector.exec()
